

package com.hazelcast.jet.impl.pipeline.transform;

import com.hazelcast.function.FunctionEx;
import com.hazelcast.jet.Traverser;
import com.hazelcast.jet.core.Edge;
import com.hazelcast.jet.impl.pipeline.PipelineImpl.Context;
import com.hazelcast.jet.impl.pipeline.Planner;
import com.hazelcast.jet.impl.pipeline.Planner.PlannerVertex;

import javax.annotation.Nonnull;
import java.io.Serial;

import static com.hazelcast.jet.core.Vertex.LOCAL_PARALLELISM_USE_DEFAULT;
import static com.hazelcast.jet.core.processor.Processors.flatMapP;

public class FlatMapTransform<T, R> extends AbstractTransform {

    @Serial
    private static final long serialVersionUID = 1L;

    @Nonnull
    private FunctionEx<? super T, ? extends Traverser<? extends R>> flatMapFn;

    public FlatMapTransform(@Nonnull String name, @Nonnull Transform upstream, @Nonnull FunctionEx<? super T, ? extends Traverser<? extends R>> flatMapFn) {
        super(name, upstream);
        this.flatMapFn = flatMapFn;
    }

    @Nonnull
    public FunctionEx<? super T, ? extends Traverser<? extends R>> flatMapFn() {
        return flatMapFn;
    }

    @Override
    public void addToDag(Planner p, Context context) {
        determineLocalParallelism(LOCAL_PARALLELISM_USE_DEFAULT, context, p.isPreserveOrder());
        PlannerVertex pv = p.addVertex(this, name(), determinedLocalParallelism(), flatMapP(flatMapFn()));
        if (p.isPreserveOrder()) {
            p.addEdges(this, pv.v, Edge::isolated);
        } else {
            p.addEdges(this, pv.v);
        }
    }
}
